#!/bin/bash

# exit script on any error
set -e

if [ "$EUID" -ne 0 ]
then
    echo "Must run as root or with sudo."
    exit 1
fi

DEVICE=${1}
TARBALL=${2}
CONFIG=${3}
DRIVESIZE=${4}

DEVICE_BOOT=${DEVICE}1
DEVICE_ROOT=${DEVICE}2

BOOT_MNT=/mnt/boot
ROOT_MNT=/mnt/root

if [[ -z "$TARBALL" ]]
then
    echo "Usage: ./mkeosdrive /dev/yourUSBdrive system_image.tar <optional 3rd parameter: /path/to/configfile.tar.gz>\
    <optional 4th parameter: desired end cap size of root partition in GB>"
    echo "If you want to provide a root partition end cap size but no config, use \"-\" without the quotes as the third paramenter."
    exit 1
fi

mkdir -p $BOOT_MNT
mkdir -p $ROOT_MNT

# Unmount all partitions on device
umount ${DEVICE}? && echo "Unmounted all partitions on $DEVICE" || echo "no filesystem mounted from $DEVICE"

# Unmount folders only if already mounted
grep -qs "$BOOT_MNT" /proc/mounts && umount $BOOT_MNT
grep -qs "$ROOT_MNT" /proc/mounts && umount $ROOT_MNT

DIR=`mktemp -d`

function cleanup {
    if [ -d "$DIR" ]
    then
        rm -rf "$DIR"
    fi
    if [ -d "$BOOT_MNT" ]
    then
        rm -rf "$BOOT_MNT"
    fi
    if [ -d "$ROOT_MNT" ]
    then
        rm -rf "$ROOT_MNT"
    fi
}
trap cleanup EXIT

echo "Press any key to continue formatting $DEVICE"
read -n1 -rsp $'or hit Ctrl+C to exit now...\n'

# deleting existing partiton table first
dd if=/dev/zero of=$DEVICE bs=512 count=1 conv=notrunc

echo "Waiting for disk sync..."
sync

parted --script $DEVICE mktable msdos
parted --script $DEVICE mkpart primary fat32 1 150MB

# If no drive size was supplied as a parameter, default to full size root partition
if [ ! -z $DRIVESIZE ]
then
    echo "Creating root partition with an end cap size of $DRIVESIZE GB"
    let "SIZE = $DRIVESIZE * 1024"
    parted --script $DEVICE mkpart primary ext3 150MB ${SIZE}MB
else
    parted --script $DEVICE mkpart primary ext3 150MB 100%
fi

echo "Waiting for disk sync..."
sync

echo "Formatting $DEVICE partitions..."
mkfs.msdos -F 32 $DEVICE_BOOT
mkfs.ext3 -q $DEVICE_ROOT

echo "Waiting for disk sync..."
sync

mount $DEVICE_BOOT $BOOT_MNT
mount $DEVICE_ROOT $ROOT_MNT

echo "Unpacking EdgeOS release image"
tar xf $TARBALL -C $DIR

####### Kernel
echo "Verifying EdgeOS kernel"
if [ `md5sum $DIR/vmlinux.tmp | awk -F ' ' '{print $1}'` != \
     `cat $DIR/vmlinux.tmp.md5` ]; then
    echo "Kernel in image is corrupted! Check your image and start over."
    exit 1
fi

echo "Copying EdgeOS kernel to boot partition"
cp $DIR/vmlinux.tmp $BOOT_MNT/vmlinux.64
cp $DIR/vmlinux.tmp.md5 $BOOT_MNT/vmlinux.64.md5

####### System
echo "Verifying EdgeOS system image"
if [ `md5sum $DIR/squashfs.tmp | awk -F ' ' '{print $1}'` != \
     `cat $DIR/squashfs.tmp.md5` ]; then
    echo "System in image is corrupted! Check your image and start over."
    exit 1
fi

echo "Copying EdgeOS system image to root partition"
mv $DIR/squashfs.tmp $ROOT_MNT/squashfs.img
mv $DIR/squashfs.tmp.md5 $ROOT_MNT/squashfs.img.md5

echo "Copying version file to the root partition"
mv $DIR/version.tmp $ROOT_MNT/version

# Writable data dir
echo "Creating EdgeOS writable data directories"
mkdir $ROOT_MNT/w
mkdir $ROOT_MNT/www

chown 33:0 $ROOT_MNT/www

####### Extract config
if [ ! -z $CONFIG ] && [ -f $CONFIG ]
then
    echo "Config archive $CONFIG was provided, extracting it to the root partition..."
    tar xf $CONFIG -C $ROOT_MNT/w
else
    echo "No config archive was provided, default settings will be used."
fi

echo "Waiting for disk sync..."
sync

echo "Unmounting filesystems on $DEVICE..."
umount $ROOT_MNT
umount $BOOT_MNT

echo "Done. You can now plug the USB drive back into your router and boot it."
